import torch
import torch.nn.functional as F
# from pytorch3d.ops import sample_farthest_points
import numpy as np

def fps(points, num_points=1024, rand_sample=True):
    # return points
    if points.shape[0] == 0:
        return np.zeros((0, 3)), np.zeros((0,), dtype=np.int64)
    
    if num_points >= points.shape[0]:
        return points, np.arange(points.shape[0], dtype=np.int64)
    
    num_points = min(num_points, points.shape[0])
    
    if rand_sample:
        indices = np.random.choice(points.shape[0], num_points, replace=False)
        sampled_points = points[indices]
        return sampled_points, indices
    else:
        if isinstance(points, np.ndarray):
            points_tensor = torch.from_numpy(points).float().unsqueeze(0)
        else:
            points_tensor = points.float().unsqueeze(0) if points.dim() == 2 else points.float()
        
        sampled_points_tensor, indices_tensor = sample_farthest_points(points_tensor, K=num_points)
        
        sampled_points = sampled_points_tensor.squeeze(0).cpu().numpy()
        indices = indices_tensor.squeeze(0).cpu().numpy()
        
        return sampled_points, indices
